package com.rmd.ai.grpc.service;

import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.pool2.BasePooledObjectFactory;
import org.apache.commons.pool2.ObjectPool;
import org.apache.commons.pool2.PooledObject;
import org.apache.commons.pool2.impl.DefaultPooledObject;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;

@Slf4j
@Component
public class AbilityGrpcChannelStatePool {
    @Value("${grpc.client.ability.server}")
    private String server;
    @Value("${grpc.client.ability.max-total}")
    private Integer poolMaxTotal;
    private ObjectPool<ManagedChannel> connectionPool;
    /**
     * 创建连接池并配置相关属性
     */
    @PostConstruct
    public void init() {
        if (StringUtils.isBlank(server)){
            return;
        }
        // 构建连接池工厂类
        BasePooledObjectFactory<ManagedChannel> factory = new BasePooledObjectFactory<ManagedChannel>() {

            @Override
            public ManagedChannel create() throws Exception {
                String[] ips = server.split(":");
                // 创建GRPC客户端连接对象
                return ManagedChannelBuilder.forAddress(ips[0].trim(), Integer.parseInt(ips[1].trim()))
                        .usePlaintext()
                        .build();
            }

            @Override
            public PooledObject<ManagedChannel> wrap(ManagedChannel channel) {
                return new DefaultPooledObject<>(channel);
            }

            @Override
            public void destroyObject(PooledObject<ManagedChannel> p) throws Exception{
                // 关闭连接对象
                p.getObject().shutdown();
            }

            @Override
            public boolean validateObject(PooledObject<ManagedChannel> p) {
                // 检查连接对象是否可用
                return !p.getObject().isShutdown();
            }
        };
        final GenericObjectPoolConfig objectPoolConfig = new GenericObjectPoolConfig();
        // 连接池的配置
        GenericObjectPoolConfig poolConfig = new GenericObjectPoolConfig();
        // 池中的最大连接数
        poolConfig.setMaxTotal(poolMaxTotal == null ? 10 : poolMaxTotal);
        // 最少的空闲连接数
        poolConfig.setMinIdle(1);
        // 最多的空闲连接数
        poolConfig.setMaxIdle(4);
        // 当连接池资源耗尽时,调用者最大阻塞的时间,超时时抛出异常 单位:毫秒数
        poolConfig.setMaxWaitMillis(5000);
        // 连接池存放池化对象方式,true放在空闲队列最前面,false放在空闲队列最后
        poolConfig.setLifo(true);
        // 连接空闲的最小时间,达到此值后空闲连接可能会被移除,默认即为30分钟
        poolConfig.setMinEvictableIdleTimeMillis(1000L * 60L * 30L);
        // 连接耗尽时是否阻塞,默认为true
        poolConfig.setBlockWhenExhausted(true);
        this.connectionPool = new GenericObjectPool<>(factory, objectPoolConfig);
    }

    public ManagedChannel borrowConnect() throws Exception {
        ManagedChannel channel = connectionPool.borrowObject();
        log.info("AbilityGrpcChannelStatePool total threads activate: " + connectionPool.getNumActive());
        // 从连接池中获取连接对象
        return channel;
    }

    public void returnConnect(ManagedChannel channel) {
        // 将连接对象归还到连接池中
        try {
            connectionPool.returnObject(channel);
        }catch (Exception e){
            log.error("AbilityGrpcChannelStatePool returnConnect",e);
        }
    }

    /**
     * 关闭连接池
     */
    public void close(){
        try {
            connectionPool.close();
        }catch (Exception e){
            log.error("AbilityGrpcChannelStatePool close",e);
        }
    }
}
